# 卷积神经网络验证码版本
import tensorflow as tf
import random
import os
import numpy as np
from PIL import Image
import sys

script_name = os.path.basename(__file__)
ver = 'v1.0'

file_path = r'../../../../large_data/DL1/vcode_data/'
items = os.listdir(file_path + 'train')
item = items[0]
print(f'item: {item}')
label = os.path.splitext(os.path.split(item)[1])[0]
print(label)
split_result = os.path.split(item)
print(f'split: {split_result}')
print(f'splitext: {os.path.splitext(split_result[1])}')  # ATTENTION os path split ext
# sys.exit(0)

np.random.seed(777)
tf.set_random_seed(777) #设置随机种子

# 获取数据集
train_num = 1000 #训练集数量
test_num = 100 #测试集数量

IMG_HEIGHT = 60 #图片的高度
IMG_WIDTH = 160 #图片的宽度
char_num = 4  # 验证码字符个数
characters = range(10)
labellen = char_num * len(characters)
def label2vec(label):
    """
    :param label: 1327
        [0,1,0,0,0,0,0,0,0,0,
        0,0,0,1,0,0,0,0,0,0,
        0,0,1,0,0,0,0,0,0,0,
        0,0,0,0,0,0,0,1,0,0]    """
    label_vec = np.zeros(char_num * len(characters))
    for i, num in enumerate(label):
        idx = i * len(characters) + int(num)
        label_vec[idx] = 1
    return label_vec
def convert2gray(img):  #彩色转成灰度
    if len(img.shape)>2:
        r, g, b = img[:, :, 0], img[:, :, 1], img[:, :, 2]
        gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
        return gray
    else:
        return img
def get_all_files(file_path, num): # 获取图片路径及其标签 file_path: a sting, 图片所在目录
    image_list = []
    label_list = []
    i = 0
    for item in os.listdir(file_path): # 文件名形如0000.jpg
        item_path = file_path + '\\' + item
        image = Image.open(item_path)
        image = convert2gray(np.array(image))
        image_array = np.array(image) / 255.0
        image_list.append(image_array.reshape(IMG_HEIGHT,IMG_WIDTH,1))
        label = os.path.splitext(os.path.split(item)[1])[0]
        #item[0:4]
        label_list.append(label2vec(label))
        i += 1
        if i >= num:
            break
    return np.array(image_list), np.array(label_list)
image_dir = file_path + 'train'
test_dir = file_path + 'test'
imgArr, Y_one_hot = get_all_files(image_dir, train_num)
imgArrTest, Y_test = get_all_files(test_dir, test_num)
print(imgArr.shape, len(Y_one_hot))
g_b=0
# 自己实现next_batch函数，每次返回一批数据
def next_batch(size):
    global g_b
    xb = imgArr[g_b:g_b+size]
    yb = Y_one_hot[g_b:g_b+size]
    g_b = g_b + size
    return xb,yb

# 参数
learning_rate = 0.001 # 学习率
training_epochs = 100  # 训练总周期
batch_size = 100 # 训练每批样本数

#定义占位符
X = tf.placeholder(tf.float32, [None, IMG_HEIGHT, IMG_WIDTH, 1])
Y = tf.placeholder(tf.float32, [None, labellen])  # 多个数字的独热编码组合

with tf.variable_scope('conv1'): # 第1层卷积，输入图片数据(?, IMG_HEIGHT, IMG_WIDTH, 1)
    W1 = tf.Variable(tf.random_normal([3, 3, 1, 32]))  #卷积核3x3，输入通道1，输出通道32
    L1 = tf.nn.conv2d(X, W1, strides=[1, 1, 1, 1], padding='SAME')
    L1 = tf.nn.relu(L1)
    L1 = tf.nn.max_pool(L1, ksize=[1, 2, 2, 1],strides=[1, 2, 2, 1], padding='SAME')

with tf.variable_scope('conv2'): # 第2层卷积
    W2 = tf.Variable(tf.random_normal([3, 3, 32, 64], stddev=0.01)) #卷积核3x3，输入通道32，输出通道64
    L2 = tf.nn.conv2d(L1, W2, strides=[1, 1, 1, 1], padding='SAME')
    L2 = tf.nn.relu(L2)
    L2 = tf.nn.max_pool(L2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
with tf.variable_scope('fc1'): #全连接1
    dim = L2.get_shape()[1].value * L2.get_shape()[2].value * L2.get_shape()[3].value
    L2_flat = tf.reshape(L2, [-1, dim])
    W4 = tf.get_variable(name="W4", shape=[dim, 1024], dtype=tf.float32, initializer=tf.contrib.layers.xavier_initializer())
    b4 = tf.Variable(tf.random_normal([1024]))
    L4 = tf.nn.relu(tf.matmul(L2_flat, W4) + b4, name="fc1")
with tf.variable_scope('softmax'): #softmax
    W5 = tf.get_variable(name="W5", shape=[1024, labellen], dtype=tf.float32, initializer=tf.contrib.layers.xavier_initializer())
    b5 = tf.Variable(tf.random_normal([labellen]))
    logits = tf.add(tf.matmul(L4, W5), b5)

# 代价函数和优化器
# sigmod_cross适用于每个类别相互独立但不互斥，如图中可以有字母和数字
# softmax_cross适用于每个类别独立且排斥的情况，如数字和字母不可以同时出现
cost = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=logits, labels=Y))
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(cost) # 优化器

# 测试模型计算准确率
predict = tf.reshape(logits, [-1, 4, 10])
max_idx_p = tf.argmax(predict, 2)

max_idx_y = tf.argmax(tf.reshape(Y, [-1, 4, 10]), 2)
correct_pred = tf.equal(max_idx_p, max_idx_y)
accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))

tf.summary.scalar('cost', cost)
tf.summary.scalar('accuracy', accuracy)
summary = tf.summary.merge_all()

# 创建会话
sess = tf.Session()
sess.run(tf.global_variables_initializer()) #全局变量初始化
fw = tf.summary.FileWriter('./_log/' + script_name + '_' + ver, sess.graph)
# 迭代训练
print('开始学习...')
g_step = 0
for epoch in range(training_epochs):
    avg_cost = 0
    avg_acc = 0
    total_batch = int(train_num / batch_size)  # 批次
    g_b = 0
    group = int(np.ceil(total_batch / 10))
    for i in range(total_batch):
        batch_xs, batch_ys = next_batch(batch_size)
        feed_dict = {X: batch_xs, Y: batch_ys}
        c, _, accv, sv = sess.run([cost, optimizer, accuracy, summary], feed_dict=feed_dict)
        avg_cost += c / total_batch
        avg_acc += accv / total_batch
        fw.add_summary(sv, g_step)
        g_step += 1
        if i % group == 0:
            print(f'g_step#{g_step + 1}: epoch#{epoch + 1}: batch#{i + 1}: cost = {c}, acc = {accv}')
            fw.flush()
    if i % group != 0:
        print(f'g_step#{g_step + 1}: epoch#{epoch + 1}: batch#{i + 1}: cost = {c}, acc = {accv}')
        fw.flush()
    print('##### Epoch:', (epoch + 1), 'cost =', avg_cost, 'acc=',sess.run(accuracy,feed_dict={X: imgArrTest, Y: Y_test}))
print('学习完成')

# 测试模型检查准确率
print('正确率:', sess.run(accuracy, feed_dict={X: imgArrTest, Y: Y_test}))

# 在测试集中随机抽一个样本进行测试
r = int(random.randint(0, test_num - 1))
print("标签: ", sess.run(tf.argmax(Y_test[r:r + 1].reshape(4, 10), 1)))
pre = sess.run([tf.argmax(tf.reshape(logits, [4, 10]), 1)], feed_dict={X: imgArrTest[r:r + 1]})
print("预测: ", pre)
